import os
import numpy as np
import nibabel as nib
from scipy.ndimage import map_coordinates, binary_erosion, generate_binary_structure
from scipy.spatial.distance import cdist
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import traceback
from typing import Tuple, Optional

plt.switch_backend('Agg')

class ResamplingImageAmount:
    """
    一个用于从CT扫描中生成三联可视化图像集的类。
    1. 为每个采样点的所有7个倾斜角度都生成完整的三联可视化图像。
    2. 仅当'Image'图的小采样区域超出CT体素边界时，才放弃该样本，避免图像出现填充值。'Location'图允许填充。
    """
    def __init__(self, ct_nifti_path: str, masks_dir: str, neighborhood_radius_voxels: int = 3):
        print("--- ResamplingImageAmount 初始化 (邻域搜索 + 倾斜采样) ---")
        if not os.path.exists(ct_nifti_path): raise FileNotFoundError(f"CT NIfT文件不存在: {ct_nifti_path}")
        self.nii_image = nib.load(ct_nifti_path)
        self.volume_data = self.nii_image.get_fdata()
        self.affine = self.nii_image.affine
        self.inv_affine = np.linalg.inv(self.affine)
        self.voxel_spacing = np.abs(self.nii_image.header.get_zooms()[:3])
        volume_physical_dims = np.array(self.volume_data.shape) * self.voxel_spacing
        self.max_volume_dimension_mm = np.max(volume_physical_dims) * 1.1
        
        print("\n--- 加载器官掩码 ---")
        self.path_organs_mask = self._load_and_combine_masks(masks_dir, ['stomach', 'duodenum'], "“观测平台”")
        if self.path_organs_mask is None: raise RuntimeError("未能加载任何胃或十二指肠掩码。")
        self.target_organs_mask = self._load_and_combine_masks(masks_dir, ['liver', 'spleen', 'pancreas', 'kidney_left', 'kidney_right', 'gallbladder'], "“观测目标”")
        if self.target_organs_mask is None: raise RuntimeError("未能加载任何目标器官掩码。")

        print(f"\n--- 通过邻域搜索寻找有效采样点并计算方向 (邻域半径: {neighborhood_radius_voxels} 体素) ---")
        self.sampling_points_voxels, self.sampling_directions_world = self._find_valid_points_and_directions_by_neighborhood(neighborhood_radius_voxels)
        if len(self.sampling_points_voxels) == 0: raise RuntimeError("未能找到任何有效的邻近采样点。")

        print(f"\n--- 初始化完成，找到 {len(self.sampling_points_voxels)} 个有效采样点 ---")

    def _load_and_combine_masks(self, masks_dir: str, organ_names: list, group_name: str) -> np.ndarray:
        print(f"  正在加载 {group_name} 组:")
        combined_mask = None; found_any = False
        for organ in organ_names:
            mask_path = os.path.join(masks_dir, f"{organ}.nii.gz")
            if os.path.exists(mask_path):
                found_any = True; print(f"    - 找到并加载: {organ}.nii.gz")
                mask_data = nib.load(mask_path).get_fdata().astype(np.bool_)
                if combined_mask is None: combined_mask = mask_data
                else: combined_mask = np.logical_or(combined_mask, mask_data)
        return combined_mask.astype(np.uint8) if found_any else None

    def _find_valid_points_and_directions_by_neighborhood(self, radius: int) -> Tuple[np.ndarray, np.ndarray]:
        """使用邻域搜索法验证壁点，并计算精确的扫描方向。"""
        print("    - 正在提取胃/十二指肠的壁...")
        struct = generate_binary_structure(3, 1)
        eroded_mask = binary_erosion(self.path_organs_mask, structure=struct, iterations=1)
        wall_mask = self.path_organs_mask & ~eroded_mask
        wall_points_voxels = np.argwhere(wall_mask > 0)
        print(f"      找到 {len(wall_points_voxels)} 个候选壁点。")

        valid_points = []; valid_directions_world = []
        
        print(f"    - 正在为每个壁点检查 {2*radius+1}x{2*radius+1}x{2*radius+1} 的邻域...")
        for i, point in enumerate(wall_points_voxels):
            if i % 20000 == 0 and i > 0: print(f"      已处理 {i} / {len(wall_points_voxels)} 个壁点...")

            min_bounds = np.maximum(0, point - radius)
            max_bounds = np.minimum(self.volume_data.shape, point + radius + 1)
            
            neighborhood = self.target_organs_mask[min_bounds[0]:max_bounds[0], min_bounds[1]:max_bounds[1], min_bounds[2]:max_bounds[2]]
            
            if np.any(neighborhood):
                relative_target_coords = np.argwhere(neighborhood)
                absolute_target_coords = relative_target_coords + min_bounds
                distances = cdist(point.reshape(1, -1), absolute_target_coords)
                closest_target_point = absolute_target_coords[np.argmin(distances)]
                
                direction_voxel = closest_target_point - point
                norm = np.linalg.norm(direction_voxel)
                if norm < 1e-6: continue
                
                affine_rot = self.affine[:3, :3]
                direction_world = affine_rot @ direction_voxel
                direction_world /= np.linalg.norm(direction_world)
                
                valid_points.append(point)
                valid_directions_world.append(direction_world)
        
        print(f"      邻域搜索完成，共找到 {len(valid_points)} 个有效采样点。")
        return np.array(valid_points), np.array(valid_directions_world)
    
    def _rotate_vector(self, vector: np.ndarray, axis: np.ndarray, angle_deg: float) -> np.ndarray:
        """使用罗德里格旋转公式围绕任意轴旋转一个向量。"""
        angle_rad = np.deg2rad(angle_deg)
        cos_a = np.cos(angle_rad); sin_a = np.sin(angle_rad)
        return (vector * cos_a +
                np.cross(axis, vector) * sin_a +
                axis * np.dot(axis, vector) * (1 - cos_a))
    
    def _is_plane_within_bounds(self, center_mm, u_vec, v_vec, plane_width_mm, plane_height_mm) -> bool:
        """检查一个采样平面（由其四个角点定义）是否完全位于CT体素数据边界之内。"""
        half_w = plane_width_mm / 2.0
        half_h = plane_height_mm / 2.0
        
        corners_mm = np.array([
            center_mm - half_w * u_vec + half_h * v_vec,
            center_mm + half_w * u_vec + half_h * v_vec,
            center_mm + half_w * u_vec - half_h * v_vec,
            center_mm - half_w * u_vec - half_h * v_vec
        ])
        
        homogeneous_coords = np.vstack([corners_mm.T, np.ones(4)])
        voxel_coords = (self.inv_affine @ homogeneous_coords)[:3, :]
        
        min_coords = np.min(voxel_coords, axis=1)
        max_coords = np.max(voxel_coords, axis=1)
        
        volume_shape_max_indices = np.array(self.volume_data.shape) - 1
        
        if np.any(min_coords < 0) or np.any(max_coords > volume_shape_max_indices):
            return False
        return True

    def _resample_plane(self, center_mm, u_vec, v_vec, plane_width_mm, plane_height_mm, resolution) -> np.ndarray:
        """
        从三维体数据中提取任意二维斜切面。边界填充仅用于'Location'图。
        """
        s_coords = np.linspace(-plane_width_mm / 2, plane_width_mm / 2, resolution[1])
        t_coords = np.linspace(-plane_height_mm / 2, plane_height_mm / 2, resolution[0])
        s_grid, t_grid = np.meshgrid(s_coords, t_coords)
        points_mm = center_mm + s_grid[..., np.newaxis] * u_vec + t_grid[..., np.newaxis] * v_vec
        
        points_flat = points_mm.reshape(-1, 3)
        homogeneous_coords = np.vstack([points_flat.T, np.ones(points_flat.shape[0])])
        voxel_coords = self.inv_affine @ homogeneous_coords
        
        resampled_data = map_coordinates(
            self.volume_data, voxel_coords[:3, :], order=1, mode='constant', cval=-1024)
        
        return resampled_data.reshape(resolution)

    def generate_triplet_visualizations(
        self, 
        level_output_dir: str, 
        location_output_dir: str, 
        image_output_dir: str, 
        metadata_output_path: str,
        num_samples: int = 100, 
        square_side_length_mm: float = 170.0,
        window_level: int = 40,
        window_width: int = 400
    ):
        """生成所有7个角度的三联图，并仅当'Image'图区域超出边界时跳过样本。"""
        print(f"\n--- 开始为 {num_samples} 个采样点生成坐标和可视化图像 ---")
        os.makedirs(level_output_dir, exist_ok=True); os.makedirs(location_output_dir, exist_ok=True); os.makedirs(image_output_dir, exist_ok=True)
        num_path_points = len(self.sampling_points_voxels)
        if num_path_points == 0: print("警告: 没有有效的采样点可供生成图像。"); return
        if num_path_points < num_samples: 
            print(f"警告: 请求生成 {num_samples} 个样本, 但只找到 {num_path_points} 个有效点。将生成所有有效点。")
            num_samples = num_path_points
        indices_to_sample = np.random.choice(num_path_points, num_samples, replace=False)
        vmin = window_level - window_width / 2; vmax = window_level + window_width / 2
        output_resolution = (256, 256)
        
        TILT_ANGLES_DEG = [0, 30, -30, 45, -45, 60, -60]
        
        valid_coords_count = 0
        print(f"坐标元数据将保存到: {metadata_output_path}")
        with open(metadata_output_path, 'w') as f_meta:
            f_meta.write("# 格式: 采样点(x,y,z) 角点1(x,y,z) 角点2(x,y,z) 角点3(x,y,z) 角点4(x,y,z)\n")
            
            for i, point_idx in enumerate(indices_to_sample):
                center_voxel_coords = self.sampling_points_voxels[point_idx]
                center_mm = nib.affines.apply_affine(self.affine, center_voxel_coords)
                direction_world_base = self.sampling_directions_world[point_idx]
                
                print(f"  - [样本 {i+1}/{num_samples}] 采样点: {center_voxel_coords} -> {np.round(center_mm, 2)} mm")
                
                v_vec_base = direction_world_base
                temp_vec = np.array([0, 0, 1]) if np.abs(np.dot(v_vec_base, [0, 0, 1])) < 0.9 else np.array([0, 1, 0])
                u_vec_base = np.cross(v_vec_base, temp_vec)
                u_vec_base /= np.linalg.norm(u_vec_base)
                
                for angle_deg in TILT_ANGLES_DEG:
                    u_vec_current = u_vec_base
                    
                    if angle_deg == 0:
                        v_vec_current = v_vec_base
                    else:
                        v_vec_current = self._rotate_vector(v_vec_base, axis=u_vec_base, angle_deg=angle_deg)
                    
                    shift_to_center = v_vec_current * (square_side_length_mm / 2.0)
                    square_center = center_mm + shift_to_center
                    
                    # 边界检查仅针对 Image 图的小正方形区域
                    is_valid = self._is_plane_within_bounds(square_center, u_vec_current, v_vec_current, square_side_length_mm, square_side_length_mm)
                    
                    if not is_valid:
                        print(f"    -> [倾斜面 {angle_deg}°] 'Image'图采样区域超出边界，已跳过。")
                        continue
                    else:
                        print(f"    -> [倾斜面 {angle_deg}°] 正在处理...")

                    # 只有在确认采样有效后，才计算角点并写入文件
                    half_side = square_side_length_mm / 2.0
                    corner1 = square_center - half_side * u_vec_current + half_side * v_vec_current
                    corner2 = square_center + half_side * u_vec_current + half_side * v_vec_current
                    corner3 = square_center + half_side * u_vec_current - half_side * v_vec_current
                    corner4 = square_center - half_side * u_vec_current - half_side * v_vec_current
                    
                    def format_coord(coord): return f"({coord[0]:.4f},{coord[1]:.4f},{coord[2]:.4f})"
                    line_to_write = f"{format_coord(center_mm)} {format_coord(corner1)} {format_coord(corner2)} {format_coord(corner3)} {format_coord(corner4)}\n"
                    f_meta.write(line_to_write)
                    valid_coords_count += 1
                    
                    base_filename = f"sample_{i:04d}_angle_{angle_deg:+d}.png"
                    
                    # 生成 Image 图
                    square_image_data = self._resample_plane(square_center, u_vec_current, v_vec_current, square_side_length_mm, square_side_length_mm, output_resolution)
                    plt.imsave(os.path.join(image_output_dir, base_filename), square_image_data, cmap='gray', vmin=vmin, vmax=vmax, origin='upper')

                    # 生成 Level 图
                    fig, ax = plt.subplots(figsize=(8, 8), facecolor='black')
                    axial_slice_idx = np.clip(int(round(center_voxel_coords[2])), 0, self.volume_data.shape[2] - 1)
                    ax.imshow(np.rot90(self.volume_data[:, :, axial_slice_idx]), cmap='gray', vmin=vmin, vmax=vmax, origin='upper')
                    plot_x, plot_y = center_voxel_coords[1], center_voxel_coords[0]
                    ax.scatter(plot_x, plot_y, c='cyan', s=60, marker='+')
                    dir_voxel_space = (self.inv_affine[:3, :3] @ v_vec_current)
                    ax.arrow(plot_x, plot_y, dir_voxel_space[1]*30, -dir_voxel_space[0]*30, color='cyan', head_width=8, width=2)
                    ax.axis('off'); plt.title(f"Angle: {angle_deg} deg", color='white');
                    plt.savefig(os.path.join(level_output_dir, base_filename), bbox_inches='tight', pad_inches=0, facecolor='black'); plt.close(fig)

                    # 生成 Location 图
                    location_plane_size_mm = self.max_volume_dimension_mm
                    location_bg_image = self._resample_plane(square_center, u_vec_current, v_vec_current, location_plane_size_mm, location_plane_size_mm, (512, 512))
                    
                    fig, ax = plt.subplots(figsize=(8, 8), facecolor='black')
                    ax.imshow(location_bg_image, cmap='gray', vmin=vmin, vmax=vmax, origin='upper')
                    mm_per_pixel = location_plane_size_mm / 512; side_length_px = square_side_length_mm / mm_per_pixel
                    rect = Rectangle((256 - side_length_px/2, 256 - side_length_px/2), side_length_px, side_length_px, facecolor='red', alpha=0.3, edgecolor='red')
                    ax.add_patch(rect)
                    ax.scatter(256, 256 - side_length_px/2, c='cyan', s=60, marker='+')
                    ax.axis('off'); plt.title(f"Angle: {angle_deg} deg", color='white')
                    plt.savefig(os.path.join(location_output_dir, base_filename), bbox_inches='tight', pad_inches=0, facecolor='black'); plt.close(fig)

        print(f"\n🎉 成功！共生成了 {valid_coords_count} 组有效坐标和图像。")
        print(f"坐标元数据已成功写入 {metadata_output_path}")



if __name__ == '__main__':
    try:
        BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        #第一步dcm2nii.gz的nii.gz文件
        CT_NIFTI_PATH = os.path.join(BASE_DIR, "data\\nifti_output\\patient001\\Abdomen-V  2.0  B30f\\HX_Chest_Abdomen_C_9.nii.gz")
        #第二步分割nii.gz得到单器官的nii.gz文件（指定胃部、十二指肠，进行采样）
        SEGMENTATION_MASKS_DIR = os.path.join(BASE_DIR, "data\\segmentation_results\\patient001\\Abdomen-V  2.0  B30f")
        # 生成的三联图像的输出目录
        DATASET_BASE_DIR = os.path.join(BASE_DIR, "data", "EUS_Triplet_Dataset_v6_All_Angles_Checked")
        # 水平切面图
        LOCATION_OUTPUT_DIR = os.path.join(DATASET_BASE_DIR, "location")
        # 采样区域图像
        IMAGE_OUTPUT_DIR = os.path.join(DATASET_BASE_DIR, "image")
        # 采样切面视图
        LEVEL_OUTPUT_DIR = os.path.join(DATASET_BASE_DIR, "level")
        #输出采样点的世界坐标
        METADATA_OUTPUT_PATH = os.path.join(DATASET_BASE_DIR, "sampling_metadata_checked.txt") 
        
        os.makedirs(DATASET_BASE_DIR, exist_ok=True)
        
        print("--- 开始执行脚本 ---")
        # 增加num_samples，并减小square_side_length_mm以获得更多、更精细的样本
        generator = ResamplingImageAmount(
            ct_nifti_path=CT_NIFTI_PATH, 
            masks_dir=SEGMENTATION_MASKS_DIR,
            neighborhood_radius_voxels=2 
        )
        
        generator.generate_triplet_visualizations(
            level_output_dir=LEVEL_OUTPUT_DIR,
            location_output_dir=LOCATION_OUTPUT_DIR,
            image_output_dir=IMAGE_OUTPUT_DIR,
            metadata_output_path=METADATA_OUTPUT_PATH, 
            num_samples=200, # 增加样本数量
            square_side_length_mm=85.0,     #mm单位，采样区域是一采样点为上边中点的正方形。正方形的边长
            window_level=40, 
            window_width=400          
        )
        
    except Exception as e:
        traceback.print_exc()
        print(f"\n发生严重错误: {e}")